import os
import numpy as np
import re
import time


def calculate_weight(x, y, center_x, center_y, w, h):
    i = abs(x - center_x)
    j = abs(y - center_y)
    return (1 - 2 * i / w) * (1 - 2 * j / h)


def get_parent_flow_files(row, col):
    return [
        (row, col, row + 1, col + 1),
        (row - 1, col - 1, row, col),
        (row - 1, col, row, col + 1),
        (row, col - 1, row + 1, col)
    ]


def merge_flows(flow_dir, output_dir, tile_size=128, full_size=512):
    os.makedirs(output_dir, exist_ok=True)

    # 获取文件名中的最大行和列
    max_row, max_col = 0, 0
    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]
    for flow_file in flow_files:
        match = re.search(r'tile_\((\d+)_(\d+),(\d+)_(\d+)\)\.npy', flow_file)
        if match:
            x1, y1, x2, y2 = map(int, match.groups())
            max_row = max(max_row, x2)
            max_col = max(max_col, y2)

    for row in range(max_row + 1):
        for col in range(max_col + 1):
            sub_flows = {}
            sub_weights = {}
            for parent in get_parent_flow_files(row, col):
                parent_file = f"tile_({parent[0]}_{parent[1]},{parent[2]}_{parent[3]}).npy"
                parent_path = os.path.join(flow_dir, parent_file)
                if os.path.exists(parent_path):
                    load_start_time = time.time()
                    flow = np.load(parent_path)
                    load_end_time = time.time()
                    print(f"Time to load {parent_file}: {load_end_time - load_start_time:.4f} seconds")

                    center_x = full_size // 2
                    center_y = full_size // 2

                    # 父光流场四分之一区域
                    for i in range(tile_size):
                        for j in range(tile_size):
                            if parent == (row, col, row + 1, col + 1):
                                x = i + tile_size * 0
                                y = j + tile_size * 0
                            elif parent == (row - 1, col - 1, row, col):
                                x = i + tile_size * 1
                                y = j + tile_size * 1
                            elif parent == (row - 1, col, row, col + 1):
                                x = i + tile_size * 1
                                y = j + tile_size * 0
                            elif parent == (row, col - 1, row + 1, col):
                                x = i + tile_size * 0
                                y = j + tile_size * 1
                            weight = calculate_weight(x, y, center_x, center_y, full_size, full_size)
                            if (i, j) not in sub_flows:
                                sub_flows[(i, j)] = flow[x, y] * weight
                                sub_weights[(i, j)] = weight
                            else:
                                sub_flows[(i, j)] += flow[x, y] * weight
                                sub_weights[(i, j)] += weight
            end_time = time.time()
            print(f"Time for calculating sub_flows for tile ({row}, {col}): {end_time - load_end_time:.4f} seconds")

            start_time = time.time()
            # 计算加权平均并保存子光流场
            avg_flow = np.zeros((tile_size, tile_size, 2), dtype=np.float32)
            for (i, j), flow in sub_flows.items():
                avg_flow[i, j] = flow / sub_weights[(i, j)]  # 对每个点的权重进行加权平均
            output_file = os.path.join(output_dir, f"tile_({row}_{col}).npy")
            np.save(output_file, avg_flow)
            end_time = time.time()
            print(f"Time for saving avg_flow for tile ({row}, {col}): {end_time - start_time:.4f} seconds")
            print(f"Saved {output_file}")


# 示例用法
flow_dir = 'E:\\wafer52\\11867_32nm_fine_align_512_dvf(1)'  # 替换为实际输入目录
output_dir = 'E:\\wafer52\\11867_32nm_fine_align_256_dvf_repair(1)'  # 替换为实际输出目录

merge_flows(flow_dir, output_dir, tile_size=256, full_size=512)
